#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import asyncio
import websockets
import json
import numpy as np
import logging
from typing import Dict, Optional, List
import threading
import time
from concurrent.futures import ThreadPoolExecutor

# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class FunASRWebSocketServer:
    def __init__(self, host="0.0.0.0", port=8765, streaming_model="paraformer-zh-streaming"):
        self.host = host
        self.port = port
        self.streaming_model = streaming_model
        self.asr_model = None
        self.vad_model = None
        self.punc_model = None
        self.clients = {}
        self.executor = ThreadPoolExecutor(max_workers=4)
        
        # 官方推荐的流式参数
        self.chunk_size = [0, 10, 5]  # [0, 10, 5] = 600ms实时显示，300ms前瞻
        self.encoder_chunk_look_back = 4  # encoder self-attention回看chunk数
        self.decoder_chunk_look_back = 1  # decoder cross-attention回看chunk数
        
        # VAD参数
        self.vad_chunk_size = 200  # ms
        
        # 音频参数（按照官方文档）
        self.sample_rate = 16000
        self.chunk_stride = self.chunk_size[1] * 960  # 600ms = 9600 samples
        self.vad_chunk_stride = int(self.vad_chunk_size * self.sample_rate / 1000)  # 200ms
        
        self._load_models()
    
    def _load_models(self):
        """加载FunASR模型"""
        try:
            from funasr import AutoModel
            
            # 加载流式ASR模型
            logger.info(f"Loading streaming ASR model: {self.streaming_model}")
            self.asr_model = AutoModel(model=self.streaming_model)
            
            # 加载VAD模型
            logger.info("Loading VAD model: fsmn-vad")
            self.vad_model = AutoModel(model="fsmn-vad")
            
            # 加载标点符号恢复模型
            logger.info("Loading punctuation model: ct-punc")
            self.punc_model = AutoModel(model="ct-punc")
            
            logger.info("All models loaded successfully")
            
        except Exception as e:
            logger.error(f"Failed to load models: {e}")
            raise
    
    async def register_client(self, websocket, path):
        """注册新客户端"""
        client_id = id(websocket)
        self.clients[client_id] = {
            'websocket': websocket,
            'asr_cache': {},
            'vad_cache': {},
            'audio_buffer': [],
            'vad_buffer': [],
            'last_activity': time.time(),
            'session_id': None,
            'complete_text': "",  # 完整转录文本
            'current_sentence': "",  # 当前句子
            'chunk_count': 0,
            'sentence_count': 0,
        }
        logger.info(f"Client {client_id} connected")
        
        try:
            await self.handle_client(websocket, client_id)
        finally:
            await self.unregister_client(client_id)
    
    async def unregister_client(self, client_id):
        """注销客户端"""
        if client_id in self.clients:
            del self.clients[client_id]
            logger.info(f"Client {client_id} disconnected")
    
    async def handle_client(self, websocket, client_id):
        """处理客户端消息"""
        try:
            async for message in websocket:
                await self.process_message(websocket, client_id, message)
        except websockets.exceptions.ConnectionClosed:
            logger.info(f"Client {client_id} connection closed")
        except Exception as e:
            logger.error(f"Error handling client {client_id}: {e}")
    
    async def process_message(self, websocket, client_id, message):
        """处理客户端消息"""
        try:
            if isinstance(message, bytes):
                # 音频数据
                await self.process_audio_data(websocket, client_id, message)
            else:
                # JSON消息
                data = json.loads(message)
                await self.process_control_message(websocket, client_id, data)
                
        except Exception as e:
            logger.error(f"Error processing message from client {client_id}: {e}")
            await self.send_error(websocket, str(e))
    
    async def process_control_message(self, websocket, client_id, data):
        """处理控制消息"""
        msg_type = data.get('type')
        
        if msg_type == 'start':
            # 开始会话
            session_id = data.get('session_id', f"session_{client_id}_{int(time.time())}")
            self.clients[client_id]['session_id'] = session_id
            
            await self.send_response(websocket, {
                'type': 'session_started',
                'session_id': session_id,
                'status': 'ready',
                'config': {
                    'chunk_size': self.chunk_size,
                    'sample_rate': self.sample_rate,
                    'vad_chunk_size': self.vad_chunk_size
                }
            })
            logger.info(f"Session started for client {client_id}: {session_id}")
            
        elif msg_type == 'stop':
            # 结束会话
            await self.finalize_session(websocket, client_id)
            
        elif msg_type == 'ping':
            # 心跳
            await self.send_response(websocket, {'type': 'pong'})
    
    async def process_audio_data(self, websocket, client_id, audio_data):
        """处理音频数据"""
        try:
            client = self.clients[client_id]
            
            # 将音频数据转换为numpy数组
            audio_array = np.frombuffer(audio_data, dtype=np.int16).astype(np.float32) / 32768.0
            
            # 过滤无效数据
            if len(audio_array) == 0:
                return
            
            # 添加到缓冲区
            client['audio_buffer'].extend(audio_array)
            client['vad_buffer'].extend(audio_array)
            client['last_activity'] = time.time()
            
            # 处理VAD检测（200ms）
            while len(client['vad_buffer']) >= self.vad_chunk_stride:
                vad_chunk = np.array(client['vad_buffer'][:self.vad_chunk_stride])
                client['vad_buffer'] = client['vad_buffer'][self.vad_chunk_stride:]
                
                # 异步处理VAD
                asyncio.create_task(self.process_vad_chunk(websocket, client_id, vad_chunk, False))
            
            # 处理ASR识别（600ms）
            while len(client['audio_buffer']) >= self.chunk_stride:
                asr_chunk = np.array(client['audio_buffer'][:self.chunk_stride])
                client['audio_buffer'] = client['audio_buffer'][self.chunk_stride:]
                
                # 异步处理ASR
                asyncio.create_task(self.process_asr_chunk(websocket, client_id, asr_chunk, False))
                
        except Exception as e:
            logger.error(f"Error processing audio data from client {client_id}: {e}")
            # 重置客户端缓冲区以避免持续错误
            if client_id in self.clients:
                self.clients[client_id]['audio_buffer'] = []
                self.clients[client_id]['vad_buffer'] = []
    
    async def process_vad_chunk(self, websocket, client_id, vad_chunk, is_final=False):
        """处理VAD检测"""
        try:
            client = self.clients[client_id]
            
            # 在线程池中执行VAD推理
            loop = asyncio.get_event_loop()
            vad_result = await loop.run_in_executor(
                self.executor,
                self._run_vad_inference,
                vad_chunk,
                client['vad_cache'],
                is_final
            )
            
            if vad_result:
                await self.send_response(websocket, vad_result)
                
        except Exception as e:
            logger.error(f"Error in VAD processing for client {client_id}: {e}")
    
    async def process_asr_chunk(self, websocket, client_id, asr_chunk, is_final=False):
        """处理ASR识别 - 按官方流式处理方式"""
        try:
            client = self.clients[client_id]
            
            # 在线程池中执行ASR推理
            loop = asyncio.get_event_loop()
            asr_result = await loop.run_in_executor(
                self.executor,
                self._run_asr_inference,
                asr_chunk,
                client['asr_cache'],
                is_final
            )
            
            if asr_result and asr_result.get('text'):
                text = asr_result['text'].strip()
                if text:
                    # 更新文本状态
                    await self.update_text_state(websocket, client_id, text, is_final, asr_result)
                
        except Exception as e:
            logger.error(f"Error in ASR processing for client {client_id}: {e}")
    
    async def update_text_state(self, websocket, client_id, new_text: str, is_final: bool, original_result: dict):
        """更新文本状态并发送完整转录结果 - 改进版本"""
        try:
            client = self.clients[client_id]
            client['chunk_count'] += 1
            
            # 改进的新句子检测逻辑
            is_new_sentence = self._is_new_sentence_improved(client['current_sentence'], new_text)
            
            if is_new_sentence and client['current_sentence']:
                # 完成当前句子并添加到完整文本
                await self._finalize_current_sentence(client)
                
                # 开始新句子
                client['current_sentence'] = new_text
                client['sentence_count'] += 1
            else:
                # 更新当前句子（选择较长的文本）
                if len(new_text) >= len(client['current_sentence']):
                    client['current_sentence'] = new_text
            
            # 构建当前完整文本
            current_complete_text = client['complete_text']
            if client['current_sentence']:
                if current_complete_text and not current_complete_text.endswith(' '):
                    current_complete_text += " "
                current_complete_text += client['current_sentence']
            
            # 准备发送的结果
            response = {
                'type': 'transcription',
                'text': new_text,
                'complete_text': current_complete_text.strip(),
                'is_final': is_final,
                'is_new_sentence': is_new_sentence,
                'timestamp': time.time(),
                'confidence': original_result.get('confidence', 0.0),
                'chunk_count': client['chunk_count'],
                'sentence_count': client['sentence_count']
            }
            
            # 如果是最终结果，完成当前句子
            if is_final and client['current_sentence']:
                await self._finalize_current_sentence(client)
                response['complete_text'] = client['complete_text'].strip()
                response['has_punctuation'] = True
                # 重置当前句子
                client['current_sentence'] = ""
            
            await self.send_response(websocket, response)
            
        except Exception as e:
            logger.error(f"Error updating text state for client {client_id}: {e}")
    
    async def _finalize_current_sentence(self, client):
        """完成当前句子并添加标点符号"""
        try:
            if not client['current_sentence']:
                return
            
            # 应用标点符号处理
            punc_text = await self._apply_punctuation(client['current_sentence'])
            
            if punc_text and punc_text != client['current_sentence']:
                # 标点符号处理成功
                final_sentence = punc_text
            else:
                # 标点符号处理失败或无变化，手动添加标点
                final_sentence = client['current_sentence']
                if not final_sentence.endswith(('。', '！', '？', '.', '!', '?')):
                    final_sentence += '。'
            
            # 添加到完整文本
            if client['complete_text']:
                if not client['complete_text'].endswith(' '):
                    client['complete_text'] += " "
                client['complete_text'] += final_sentence
            else:
                client['complete_text'] = final_sentence
            
        except Exception as e:
            logger.error(f"Error finalizing sentence: {e}")
    
    def _is_new_sentence_improved(self, current: str, new_text: str) -> bool:
        """改进的新句子检测逻辑"""
        if not current:
            return True
        
        # 如果新文本明显比当前文本短，可能是新句子
        if len(new_text) < len(current) * 0.5:
            return True
        
        # 如果新文本与当前文本的相似度很低，可能是新句子
        similarity = self._calculate_similarity(current, new_text)
        if similarity < 0.3:
            return True
        
        # 如果新文本包含明显的句子结束标志
        if any(punct in new_text for punct in ['。', '！', '？', '.', '!', '?']):
            # 检查是否在文本中间或末尾
            for punct in ['。', '！', '？', '.', '!', '?']:
                if punct in new_text:
                    idx = new_text.rfind(punct)
                    # 如果标点符号后面还有文字，可能是新句子
                    if idx < len(new_text) - 1 and new_text[idx+1:].strip():
                        return True
        
        return False
    
    def _calculate_similarity(self, text1: str, text2: str) -> float:
        """计算两个文本的相似度"""
        if not text1 or not text2:
            return 0.0
        
        # 计算最长公共子序列
        common_length = 0
        min_len = min(len(text1), len(text2))
        
        for i in range(min_len):
            if text1[i] == text2[i]:
                common_length += 1
            else:
                break
        
        return common_length / max(len(text1), len(text2))
    
    async def _apply_punctuation(self, text: str) -> Optional[str]:
        """应用标点符号处理"""
        try:
            if not text.strip():
                return None
            
            loop = asyncio.get_event_loop()
            punc_text = await loop.run_in_executor(
                self.executor,
                self._run_punctuation_inference,
                text
            )
            return punc_text
        except Exception as e:
            logger.error(f"Error applying punctuation: {e}")
            return text
    
    def _run_vad_inference(self, vad_chunk, vad_cache, is_final):
        """运行VAD推理（在线程池中执行）"""
        try:
            # 验证输入数据
            if len(vad_chunk) == 0:
                return None
            
            # 确保VAD chunk大小正确
            expected_chunk_size = int(self.vad_chunk_size * self.sample_rate / 1000)
            
            # 调整chunk大小
            if len(vad_chunk) < expected_chunk_size and not is_final:
                padding = np.zeros(expected_chunk_size - len(vad_chunk), dtype=np.float32)
                vad_chunk = np.concatenate([vad_chunk, padding])
            elif len(vad_chunk) > expected_chunk_size:
                vad_chunk = vad_chunk[:expected_chunk_size]
            
            # 清理损坏的缓存
            if isinstance(vad_cache, dict) and 'stats' in vad_cache:
                try:
                    stats = vad_cache['stats']
                    if not isinstance(stats, (list, np.ndarray)) or len(stats) == 0:
                        logger.warning("Invalid VAD cache stats, clearing cache")
                        vad_cache.clear()
                except (KeyError, TypeError, IndexError):
                    logger.warning("Corrupted VAD cache detected, clearing")
                    vad_cache.clear()
            
            # 使用官方VAD API
            try:
                vad_result = self.vad_model.generate(
                    input=vad_chunk,
                    cache=vad_cache,
                    is_final=is_final,
                    chunk_size=self.vad_chunk_size
                )
            except (IndexError, KeyError) as e:
                logger.warning(f"VAD error detected: {e}, clearing cache and retrying")
                vad_cache.clear()
                try:
                    vad_result = self.vad_model.generate(
                        input=vad_chunk,
                        cache={},
                        is_final=is_final,
                        chunk_size=self.vad_chunk_size
                    )
                    vad_cache.clear()
                except Exception as retry_e:
                    logger.error(f"VAD retry failed: {retry_e}")
                    return None
            
            if vad_result and len(vad_result[0]["value"]) > 0:
                return {
                    'type': 'vad',
                    'segments': vad_result[0]["value"],
                    'timestamp': time.time(),
                    'is_final': is_final
                }
            
            return None
            
        except Exception as e:
            logger.error(f"Error in VAD inference: {e}")
            try:
                vad_cache.clear()
            except:
                pass
            return None
    
    def _run_asr_inference(self, asr_chunk, asr_cache, is_final):
        """运行ASR推理 - 按照官方流式实现"""
        try:
            # 验证输入数据
            if len(asr_chunk) == 0:
                return None
                
            # 确保音频数据为正确的numpy格式
            if not isinstance(asr_chunk, np.ndarray):
                asr_chunk = np.array(asr_chunk, dtype=np.float32)
            
            # 确保chunk大小正确
            expected_chunk_size = self.chunk_stride
            
            if len(asr_chunk) < expected_chunk_size and not is_final:
                padding = np.zeros(expected_chunk_size - len(asr_chunk), dtype=np.float32)
                asr_chunk = np.concatenate([asr_chunk, padding])
            
            # 检查并清理损坏的ASR缓存
            if isinstance(asr_cache, dict):
                try:
                    for key in list(asr_cache.keys()):
                        if isinstance(asr_cache[key], (list, np.ndarray)):
                            if hasattr(asr_cache[key], '__len__') and len(asr_cache[key]) == 0:
                                continue
                        elif asr_cache[key] is None:
                            continue
                except (AttributeError, TypeError):
                    logger.warning("Corrupted ASR cache detected, clearing")
                    asr_cache.clear()
            
            # 使用官方流式ASR API - 完全按照官方文档
            try:
                asr_result = self.asr_model.generate(
                    input=asr_chunk,
                    cache=asr_cache,
                    is_final=is_final,
                    chunk_size=self.chunk_size,
                    encoder_chunk_look_back=self.encoder_chunk_look_back,
                    decoder_chunk_look_back=self.decoder_chunk_look_back
                )
            except (IndexError, KeyError, RuntimeError) as e:
                logger.warning(f"ASR error detected: {e}, clearing cache and retrying")
                asr_cache.clear()
                try:
                    asr_result = self.asr_model.generate(
                        input=asr_chunk,
                        cache={},
                        is_final=is_final,
                        chunk_size=self.chunk_size,
                        encoder_chunk_look_back=self.encoder_chunk_look_back,
                        decoder_chunk_look_back=self.decoder_chunk_look_back
                    )
                    asr_cache.clear()
                except Exception as retry_e:
                    logger.error(f"ASR retry failed: {retry_e}")
                    return None
            
            if asr_result and len(asr_result[0]["text"]) > 0:
                text = asr_result[0]["text"].strip()
                
                if text:
                    return {
                        'type': 'transcription',
                        'text': text,
                        'is_final': is_final,
                        'timestamp': time.time(),
                        'confidence': asr_result[0].get("confidence", 0.0)
                    }
            
            return None
            
        except Exception as e:
            logger.error(f"Error in ASR inference: {e}")
            if "index" in str(e) and "out of bounds" in str(e):
                logger.warning("Detected index error, clearing ASR cache")
                try:
                    asr_cache.clear()
                except:
                    pass
            return None
    
    def _run_punctuation_inference(self, text):
        """运行标点符号推理"""
        try:
            if not text.strip():
                return None
            
            # 使用官方标点符号API
            punc_result = self.punc_model.generate(input=text)
            
            if punc_result and len(punc_result[0]["text"]) > 0:
                return punc_result[0]["text"]
            
            return text
            
        except Exception as e:
            logger.error(f"Error in punctuation inference: {e}")
            return text
    
    async def finalize_session(self, websocket, client_id):
        """结束会话，处理剩余音频和文本 - 改进版本"""
        try:
            client = self.clients[client_id]
            
            # 处理剩余的VAD数据
            if client['vad_buffer'] and len(client['vad_buffer']) >= 480:
                remaining_vad = np.array(client['vad_buffer'])
                await self.process_vad_chunk(websocket, client_id, remaining_vad, True)
            
            # 处理剩余的ASR数据，确保最后一段也能被处理
            if client['audio_buffer'] and len(client['audio_buffer']) >= 960:
                remaining_asr = np.array(client['audio_buffer'])
                await self.process_asr_chunk(websocket, client_id, remaining_asr, True)
            
            # 等待一下让最后的处理完成
            await asyncio.sleep(0.1)
            
            # 强制完成当前句子
            if client['current_sentence']:
                await self._finalize_current_sentence(client)
                
                # 发送最终完整结果
                final_result = {
                    'type': 'final_result',
                    'complete_text': client['complete_text'].strip(),
                    'session_id': client['session_id'],
                    'total_chunks': client['chunk_count'],
                    'total_sentences': client['sentence_count'],
                    'timestamp': time.time()
                }
                await self.send_response(websocket, final_result)
            
            # 清理缓冲区和缓存
            client['audio_buffer'] = []
            client['vad_buffer'] = []
            client['complete_text'] = ""
            client['current_sentence'] = ""
            client['asr_cache'] = {}
            client['vad_cache'] = {}
            client['chunk_count'] = 0
            client['sentence_count'] = 0
            
            await self.send_response(websocket, {
                'type': 'session_ended',
                'session_id': client['session_id']
            })
            
            logger.info(f"Session ended for client {client_id}")
            
        except Exception as e:
            logger.error(f"Error finalizing session for client {client_id}: {e}")
    
    async def send_response(self, websocket, data):
        """发送响应给客户端"""
        try:
            await websocket.send(json.dumps(data, ensure_ascii=False))
        except Exception as e:
            logger.error(f"Error sending response: {e}")
    
    async def send_error(self, websocket, error_message):
        """发送错误消息给客户端"""
        await self.send_response(websocket, {
            'type': 'error',
            'message': error_message,
            'timestamp': time.time()
        })
    
    async def start_server(self):
        """启动WebSocket服务器"""
        logger.info(f"Starting FunASR WebSocket server on {self.host}:{self.port}")
        logger.info(f"ASR model: {self.streaming_model}")
        logger.info(f"Chunk size: {self.chunk_size} (600ms with 300ms lookahead)")
        logger.info(f"VAD chunk size: {self.vad_chunk_size}ms")
        
        async with websockets.serve(
            self.register_client,
            self.host,
            self.port,
            ping_interval=30,
            ping_timeout=10,
            max_size=1024*1024*10  # 10MB max message size
        ):
            logger.info("Server started successfully")
            await asyncio.Future()  # 保持服务器运行

def main():
    import argparse
    
    parser = argparse.ArgumentParser(description="FunASR WebSocket Server for Real-time Speech Recognition")
    parser.add_argument("--host", default="localhost", help="Server host")
    parser.add_argument("--port", type=int, default=8765, help="Server port")
    parser.add_argument("--model", default="paraformer-zh-streaming", help="Streaming ASR model name")
    
    args = parser.parse_args()
    
    # 创建服务器
    server = FunASRWebSocketServer(
        host=args.host,
        port=args.port,
        streaming_model=args.model
    )
    
    # 启动服务器
    try:
        asyncio.run(server.start_server())
    except KeyboardInterrupt:
        logger.info("Server shutdown by user")
    except Exception as e:
        logger.error(f"Server error: {e}")

if __name__ == "__main__":
    main() 